import struct
import hashlib
import time
import sys


def print_progress_bar(iteration, total, prefix='', suffix='', length=50, fill='█', print_end="\r"):
    """
    打印进度条到控制台
    :param iteration: 当前迭代（即已完成的部分）
    :param total: 总迭代次数（即总量）
    :param prefix: 进度条前缀
    :param suffix: 进度条后缀
    :param length: 进度条长度
    :param fill: 填充字符
    :param print_end: 结尾字符
    """
    percent = ("{0:.1f}").format(100 * (iteration / float(total)))
    filled_length = int(length * iteration // total)
    bar = fill * filled_length + '-' * (length - filled_length)
    sys.stdout.write(f'\r{prefix} |{bar}| {percent}% {suffix}')
    sys.stdout.flush()


def _validate_method(method):
    methods = {'upload': 0, 'download': 1, 'show': 2, 'message': 3}
    if method not in methods:
        raise ValueError("Method must be 'upload', 'download', 'show' or 'message'")
    return methods[method]


def _generate_md5(data):
    return hashlib.md5(data).hexdigest()


def receive_full_data(socket, header_size):
    # 接收头部 # 检查头部数据是否完整 # 解析头部以获取文件名长度和数据长度
    header_data = socket.recv(header_size)
    if not header_data:
        print('错误：未接收到头部数据。')
        return None
    if len(header_data) < header_size:
        print('错误：头部数据不完整。')
        return None
    version, magic, method_id, filename_length, data_length, md5_checksum, timestamp = struct.unpack(
        MyProtocol.header_format, header_data)
    # print(f'接收到的header结构：版本：{version}，魔数：{magic}\n方法类型：{method_id},文件名长度：{filename_length}'
    #       f'数据大小：{data_length},时间戳：{timestamp},md5校验：{md5_checksum.hex()}')
    # 计算还需接收的数据长度 # 接收剩余的数据 # 更新进度条
    total_data_size = header_size + filename_length + data_length
    total_received_data = header_data
    while len(total_received_data) < total_data_size:
        more_data = socket.recv(409600 * 4)
        if not more_data:
            print('错误：在接收数据过程中连接中断。')
            return None
        total_received_data += more_data
        if method_id != 3:
            print_progress_bar(len(total_received_data),
                               total_data_size, prefix='进度:', suffix='完成', length=50)
    # 确保接收了完整的数据
    if len(total_received_data) < total_data_size:
        print('错误：数据不完整。')
        return None
    print()
    return total_received_data


class MyProtocol:
    methods = {0: 'upload', 1: 'download', 2: 'show', 3: 'message'}
    # self.version, self.magic,
    # self.method, filename_length, self.data_length, md5_bytes(16s),
    # self.timestamp
    header_format = '=BIBHI16sQ'  # 协议定义，强制对齐 36B
    header_size = struct.calcsize(header_format)  # 计算头部大小

    def __init__(self, version: int, magic: int, method: str, filename: str, data: bytes, md5_test=None):
        self.md5_checksum = _generate_md5(md5_test) if md5_test is not None else _generate_md5(data)
        self.version = version
        self.magic = magic
        self.method = _validate_method(method)
        self.filename = filename
        self.data = data
        self.data_length = len(data)
        self.timestamp = int(time.time())

    def encode(self):
        filename_encoded = self.filename.encode('utf-8')
        filename_length = len(filename_encoded)
        md5_bytes = bytes.fromhex(self.md5_checksum)
        header = struct.pack(MyProtocol.header_format, self.version, self.magic,
                             self.method, filename_length, self.data_length, md5_bytes,
                             self.timestamp)
        return header + filename_encoded + self.data

    @staticmethod
    def decode(data):
        # 调试信息：输出头部大小和数据长度,数据接收长度已经在receive 方法检验过了
        # print("Header size:", MyProtocol.header_size, "Total data length:", len(data))
        # 从数据中解析头部
        header = data[:MyProtocol.header_size]
        version, magic, method_id, filename_length, data_length, md5_checksum, timestamp = \
            struct.unpack(MyProtocol.header_format, header)
        filename_end = MyProtocol.header_size + filename_length
        # 解析文件名
        filename = data[MyProtocol.header_size:filename_end].decode('utf-8')
        method = MyProtocol.methods.get(method_id, 'unknown')
        actual_data = data[filename_end:filename_end + data_length]
        protocol_data = MyProtocol(version, magic, method, filename, actual_data)
        print(f'接收需求 version:2,magic:4660,md5:{md5_checksum.hex()}')
        print(f'接收实际 version:{protocol_data.version},magic:{protocol_data.magic},'
              f'md5:{protocol_data.md5_checksum}')
        if not protocol_data.is_valid(md5_checksum.hex()):
            return 0
        return protocol_data

    @staticmethod
    def create_completion_message(message_code):
        """
        创建一个协议消息，用于传递简单的信息，如确认消息或错误通知。
        :param message_code: 消息代码，0代表失败，1代表成功。
        :return: 编码后的协议消息。
        """
        data = struct.pack('B', message_code)  # 消息代码作为数据
        protocol = MyProtocol(
            version=2,
            magic=0x1234,
            method='message',
            filename='',
            data=data
        )
        print(f'发送：{protocol}')
        return protocol.encode()

    def __repr__(self):
        return (f"MyProtocol(version={self.version}, magic={hex(self.magic)}, "
                f"method={self.method}, data_length={self.data_length}, "
                f"fileName={self.filename},"
                f"timestamp={self.timestamp}, md5_checksum={self.md5_checksum})")

    def is_valid(self, expected_md5):
        """
        验证协议的有效性：检查MD5校验值和魔数是否一致。
        :param expected_md5: 收到的MD5
        :return: 布尔值，如果MD5和魔数和版本都正确，则为True
        """
        # 比较MD5和魔数
        return self.md5_checksum == expected_md5 and self.magic == 0x1234 and self.version == 2


# 使用示例
if __name__ == "__main__":
    try:

        protocol = MyProtocol(
            version=10,
            magic=0x1234,
            method='upload',
            filename="filename.xxx",
            data=b'1111111'
        )
        encoded_data = protocol.encode()
        decoded_protocol = MyProtocol.decode(encoded_data)
        print(decoded_protocol)
    except ValueError as e:
        print(e)
